//go:build linux
// +build linux

package collector

import (
	"encoding/csv"
	"encoding/xml"
	"fmt"
	"os/exec"
	"strconv"
	"strings"

	"sysinfo/internal/types"
	"sysinfo/internal/utils"
)

// NvidiaSMILog 表示nvidia-smi的XML输出
// 这是用于解析XML输出的简化结构
type NvidiaSMILog struct {
	GPUs []NvidiaGPU `xml:"gpu"`
}

type NvidiaGPU struct {
	ProductName string `xml:"product_name"`
	UUID        string `xml:"uuid"`
	PCIBus      string `xml:"pci>pci_bus"`
	Temperature struct {
		Current string `xml:"gpu_temp"`
	} `xml:"temperature"`
	Utilization struct {
		GPU    string `xml:"gpu_util"`
		Memory string `xml:"memory_util"`
	} `xml:"utilization"`
	FBMemory struct {
		Total string `xml:"total"`
		Used  string `xml:"used"`
		Free  string `xml:"free"`
	} `xml:"fb_memory_usage"`
	Power struct {
		Draw  string `xml:"power_draw"`
		Limit string `xml:"power_limit"`
	} `xml:"power_readings"`
	Clocks struct {
		Graphics string `xml:"graphics_clock"`
		Memory   string `xml:"mem_clock"`
	} `xml:"clocks"`
	FanSpeed      string `xml:"fan_speed"`
	DriverVersion string `xml:"driver_version"`
}

// collectGPUPlatform 实现Linux特定的GPU数据收集
func collectGPUPlatform() []types.GPUInfo {
	gpus := make([]types.GPUInfo, 0)

	// 首先尝试NVIDIA GPU（nvidia-smi）
	nvidiaGPUs := collectNvidiaGPUs()
	gpus = append(gpus, nvidiaGPUs...)

	// 尝试AMD GPU（rocm-smi或lspci）
	amdGPUs := collectAMDGPUs()
	gpus = append(gpus, amdGPUs...)

	// 如果其他方法都失败了，回退到lspci获取基本信息
	if len(gpus) == 0 {
		gpus = collectGPUsFromLspci()
	}

	return gpus
}

// collectNvidiaGPUs 使用nvidia-smi收集NVIDIA GPU信息
func collectNvidiaGPUs() []types.GPUInfo {
	gpus := make([]types.GPUInfo, 0)

	// Check if nvidia-smi is available
	_, err := exec.LookPath("nvidia-smi")
	if err != nil {
		return gpus
	}

	// Try XML format first (more detailed)
	cmd := exec.Command("nvidia-smi", "-q", "-x")
	output, err := cmd.Output()
	if err == nil {
		var smiLog NvidiaSMILog
		if err := xml.Unmarshal(output, &smiLog); err == nil {
			for i, gpu := range smiLog.GPUs {
				gpuInfo := types.GPUInfo{
					Index:         i,
					Name:          gpu.ProductName,
					Vendor:        "NVIDIA",
					Driver:        "nvidia",
					DriverVersion: gpu.DriverVersion,
					UUID:          gpu.UUID,
					PCIBus:        gpu.PCIBus,
				}

				// 解析温度
				if temp, err := strconv.Atoi(strings.TrimSpace(strings.Replace(gpu.Temperature.Current, "C", "", -1))); err == nil {
					gpuInfo.Temperature = temp
				}

				// 解析利用率
				if util, err := strconv.Atoi(strings.TrimSpace(strings.Replace(gpu.Utilization.GPU, "%", "", -1))); err == nil {
					gpuInfo.Utilization = util
				}
				if memUtil, err := strconv.Atoi(strings.TrimSpace(strings.Replace(gpu.Utilization.Memory, "%", "", -1))); err == nil {
					gpuInfo.MemoryUtilization = memUtil
				}

				// 解析内存（格式："12345 MiB"）
				if total := parseMemoryMiB(gpu.FBMemory.Total); total > 0 {
					gpuInfo.MemoryTotal = total
					gpuInfo.MemoryFormatted = utils.FormatBytes(total)
				}
				if used := parseMemoryMiB(gpu.FBMemory.Used); used > 0 {
					gpuInfo.MemoryUsed = used
				}
				if free := parseMemoryMiB(gpu.FBMemory.Free); free > 0 {
					gpuInfo.MemoryFree = free
				}

				// 解析功耗（格式："123.45 W"）
				if power := parsePowerWatts(gpu.Power.Draw); power > 0 {
					gpuInfo.PowerDraw = power
				}
				if limit := parsePowerWatts(gpu.Power.Limit); limit > 0 {
					gpuInfo.PowerLimit = limit
				}

				// 解析时钟频率（格式："1234 MHz"）
				if clock := parseClockMHz(gpu.Clocks.Graphics); clock > 0 {
					gpuInfo.ClockSpeed = clock
				}
				if memClock := parseClockMHz(gpu.Clocks.Memory); memClock > 0 {
					gpuInfo.ClockSpeedMemory = memClock
				}

				// 解析风扇转速
				if fan, err := strconv.Atoi(strings.TrimSpace(strings.Replace(gpu.FanSpeed, "%", "", -1))); err == nil {
					gpuInfo.FanSpeed = fan
				}

				gpus = append(gpus, gpuInfo)
			}
			return gpus
		}
	}

	// 回退到CSV格式获取基本信息
	cmd = exec.Command("nvidia-smi", "--query-gpu=index,name,memory.total,memory.used,memory.free,temperature.gpu,utilization.gpu,power.draw,power.limit,clocks.current.graphics,clocks.current.memory,fan.speed", "--format=csv,noheader,nounits")
	output, err = cmd.Output()
	if err != nil {
		return nil
	}

	reader := csv.NewReader(strings.NewReader(string(output)))
	records, err := reader.ReadAll()
	if err != nil {
		return gpus
	}

	for _, record := range records {
		if len(record) < 6 {
			continue
		}

		gpuInfo := types.GPUInfo{
			Vendor: "NVIDIA",
			Driver: "nvidia",
		}

		if idx, err := strconv.Atoi(strings.TrimSpace(record[0])); err == nil {
			gpuInfo.Index = idx
		}
		gpuInfo.Name = strings.TrimSpace(record[1])
		if temp, err := strconv.Atoi(strings.TrimSpace(record[2])); err == nil {
			gpuInfo.Temperature = temp
		}
		if util, err := strconv.Atoi(strings.TrimSpace(record[3])); err == nil {
			gpuInfo.Utilization = util
		}
		if total, err := strconv.ParseUint(strings.TrimSpace(record[4]), 10, 64); err == nil {
			gpuInfo.MemoryTotal = total * 1024 * 1024 // Convert MiB to bytes
			gpuInfo.MemoryFormatted = utils.FormatBytes(gpuInfo.MemoryTotal)
		}
		if used, err := strconv.ParseUint(strings.TrimSpace(record[5]), 10, 64); err == nil {
			gpuInfo.MemoryUsed = used * 1024 * 1024 // Convert MiB to bytes
		}

		gpus = append(gpus, gpuInfo)
	}

	return gpus
}

// collectAMDGPUs 使用rocm-smi收集AMD GPU信息
func collectAMDGPUs() []types.GPUInfo {
	gpus := make([]types.GPUInfo, 0)

	// 检查rocm-smi是否可用
	_, err := exec.LookPath("rocm-smi")
	if err != nil {
		return gpus
	}

	// rocm-smi --showproductname --showtemp --showuse --showmeminfo vram
	cmd := exec.Command("rocm-smi", "--showproductname", "--showtemp", "--showuse", "--showmeminfo", "vram", "--csv")
	output, err := cmd.Output()
	if err != nil {
		return gpus
	}

	// 解析CSV输出
	lines := strings.Split(string(output), "\n")
	for i, line := range lines {
		if i == 0 || strings.TrimSpace(line) == "" {
			continue
		}

		fields := strings.Split(line, ",")
		if len(fields) < 4 {
			continue
		}

		gpuInfo := types.GPUInfo{
			Index:  i - 1,
			Vendor: "AMD",
			Driver: "amdgpu",
		}

		// Parse fields (format varies, this is approximate)
		for _, field := range fields {
			field = strings.TrimSpace(field)
			if strings.Contains(strings.ToLower(field), "radeon") || strings.Contains(strings.ToLower(field), "rx") {
				gpuInfo.Name = field
			}
		}

		gpus = append(gpus, gpuInfo)
	}

	return gpus
}

// collectGPUsFromLspci 使用lspci作为获取基本GPU信息的回退方案
func collectGPUsFromLspci() []types.GPU {
	gpus := make([]types.GPUInfo, 0)

	cmd := exec.Command("lspci")
	output, err := cmd.Output()
	if err != nil {
		return gpus
	}

	lines := strings.Split(string(output), "\n")
	index := 0
	for _, line := range lines {
		lineLower := strings.ToLower(line)
		if strings.Contains(lineLower, "vga") || strings.Contains(lineLower, "3d controller") {
			gpuInfo := types.GPUInfo{
				Index: index,
			}

			// 提取PCI总线
	if parts := strings.Split(line, " "); len(parts) > 0 {
		gpuInfo.PCIBus = parts[0]
	}

	// 确定厂商和名称
	if strings.Contains(lineLower, "nvidia") {
		gpuInfo.Vendor = "NVIDIA"
		gpuInfo.Name = extractGPUName(line, "NVIDIA")
	} else if strings.Contains(lineLower, "amd") || strings.Contains(lineLower, "ati") {
		gpuInfo.Vendor = "AMD"
		gpuInfo.Name = extractGPUName(line, "AMD")
	} else if strings.Contains(lineLower, "intel") {
		gpuInfo.Vendor = "Intel"
		gpuInfo.Name = extractGPUName(line, "Intel")
	} else {
		gpuInfo.Name = line
	}

			gpus = append(gpus, gpuInfo)
			index++
		}
	}

	return gpus
}

// Helper functions

// parseMemoryMiB 解析内存大小，格式如"12345 MiB"
func parseMemoryMiB(memStr string) uint64 {
	memStr = strings.TrimSpace(memStr)
	memStr = strings.Replace(memStr, "MiB", "", -1)
	memStr = strings.TrimSpace(memStr)
	if val, err := strconv.ParseUint(memStr, 10, 64); err == nil {
		return val * 1024 * 1024 // Convert MiB to bytes
	}
	return 0
}

// parsePowerWatts 解析功耗，格式如"123.45 W"
func parsePowerWatts(powerStr string) float64 {
	powerStr = strings.TrimSpace(powerStr)
	powerStr = strings.Replace(powerStr, "W", "", -1)
	powerStr = strings.TrimSpace(powerStr)
	if val, err := strconv.ParseFloat(powerStr, 64); err == nil {
		return val
	}
	return 0
}

// parseClockMHz 解析时钟频率，格式如"1234 MHz"
func parseClockMHz(clockStr string) int {
	clockStr = strings.TrimSpace(clockStr)
	clockStr = strings.Replace(clockStr, "MHz", "", -1)
	clockStr = strings.TrimSpace(clockStr)
	if val, err := strconv.Atoi(clockStr); err == nil {
		return val
	}
	return 0
}

// extractGPUName 从lspci输出中提取GPU名称
func extractGPUName(line, vendor string) string {
	// 尝试从lspci输出中提取GPU名称
	// 常见模式："NVIDIA Corporation GP107 [GeForce GTX 1050 Ti]"
	// 或 "Advanced Micro Devices, Inc. [AMD/ATI] Ellesmere [Radeon RX 580]"
	
	// 移除厂商名称
	name := line
	if vendor == "NVIDIA" {
		name = strings.Replace(name, "NVIDIA Corporation", "", 1)
	} else if vendor == "AMD" {
		name = strings.Replace(name, "Advanced Micro Devices, Inc. [AMD/ATI]", "", 1)
		name = strings.Replace(name, "Advanced Micro Devices, Inc.", "", 1)
	} else if vendor == "Intel" {
		name = strings.Replace(name, "Intel Corporation", "", 1)
	}
	
	name = strings.TrimSpace(name)
	
	// 如果存在括号，提取其中的名称
	if strings.Contains(name, "[") && strings.Contains(name, "]") {
		start := strings.Index(name, "[")
		end := strings.Index(name, "]")
		if start < end {
			name = name[start+1 : end]
		}
	}
	
	// 清理名称
	name = strings.TrimSpace(name)
	if name == "" {
		name = "Unknown " + vendor + " GPU"
	}
	return name
}
